// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/exec/aggregate_companion_adapter.h>

#include <pollux/exec/aggregate_companion_signatures.h>
#include <pollux/exec/aggregate_function_registry.h>
#include <pollux/exec/row_container.h>
#include <pollux/expression/signature_binder.h>

namespace kumo::pollux::exec {
    void AggregateCompanionFunctionBase::setOffsetsInternal(
        int32_t offset,
        int32_t nullByte,
        uint8_t nullMask,
        int32_t initializedByte,
        uint8_t initializedMask,
        int32_t rowSizeOffset) {
        fn_->setOffsets(
            offset,
            nullByte,
            nullMask,
            initializedByte,
            initializedMask,
            rowSizeOffset);
    }

    int32_t AggregateCompanionFunctionBase::accumulatorFixedWidthSize() const {
        return fn_->accumulatorFixedWidthSize();
    }

    int32_t AggregateCompanionFunctionBase::accumulatorAlignmentSize() const {
        return fn_->accumulatorAlignmentSize();
    }

    bool AggregateCompanionFunctionBase::accumulatorUsesExternalMemory() const {
        return fn_->accumulatorUsesExternalMemory();
    }

    bool AggregateCompanionFunctionBase::isFixedSize() const {
        return fn_->isFixedSize();
    }

    void AggregateCompanionFunctionBase::setAllocatorInternal(
        HashStringAllocator *allocator) {
        fn_->setAllocator(allocator);
    }

    void AggregateCompanionFunctionBase::destroy(melon::Range<char **> groups) {
        fn_->destroy(groups);
    }

    void AggregateCompanionFunctionBase::destroyInternal(
        melon::Range<char **> groups) {
        fn_->destroy(groups);
    }

    void AggregateCompanionFunctionBase::clearInternal() {
        fn_->clear();
    }

    void AggregateCompanionFunctionBase::initializeNewGroups(
        char **groups,
        melon::Range<const vector_size_t *> indices) {
        fn_->initializeNewGroups(groups, indices);
    }

    void AggregateCompanionFunctionBase::initializeNewGroupsInternal(
        char **groups,
        melon::Range<const vector_size_t *> indices) {
        fn_->initializeNewGroups(groups, indices);
    }

    void AggregateCompanionFunctionBase::addRawInput(
        char **groups,
        const SelectivityVector &rows,
        const std::vector<VectorPtr> &args,
        bool mayPushdown) {
        fn_->addRawInput(groups, rows, args, mayPushdown);
    }

    void AggregateCompanionFunctionBase::addSingleGroupRawInput(
        char *group,
        const SelectivityVector &rows,
        const std::vector<VectorPtr> &args,
        bool mayPushdown) {
        fn_->addSingleGroupRawInput(group, rows, args, mayPushdown);
    }

    void AggregateCompanionFunctionBase::addIntermediateResults(
        char **groups,
        const SelectivityVector &rows,
        const std::vector<VectorPtr> &args,
        bool mayPushdown) {
        fn_->addIntermediateResults(groups, rows, args, mayPushdown);
    }

    void AggregateCompanionFunctionBase::addSingleGroupIntermediateResults(
        char *group,
        const SelectivityVector &rows,
        const std::vector<VectorPtr> &args,
        bool mayPushdown) {
        fn_->addSingleGroupIntermediateResults(group, rows, args, mayPushdown);
    }

    void AggregateCompanionFunctionBase::extractAccumulators(
        char **groups,
        int32_t numGroups,
        VectorPtr *result) {
        fn_->extractAccumulators(groups, numGroups, result);
    }

    void AggregateCompanionAdapter::PartialFunction::extractValues(
        char **groups,
        int32_t numGroups,
        VectorPtr *result) {
        fn_->extractAccumulators(groups, numGroups, result);
    }

    void AggregateCompanionAdapter::MergeFunction::addRawInput(
        char **groups,
        const SelectivityVector &rows,
        const std::vector<VectorPtr> &args,
        bool mayPushdown) {
        fn_->enableValidateIntermediateInputs();
        fn_->addIntermediateResults(groups, rows, args, mayPushdown);
    }

    void AggregateCompanionAdapter::MergeFunction::addSingleGroupRawInput(
        char *group,
        const SelectivityVector &rows,
        const std::vector<VectorPtr> &args,
        bool mayPushdown) {
        fn_->enableValidateIntermediateInputs();
        fn_->addSingleGroupIntermediateResults(group, rows, args, mayPushdown);
    }

    void AggregateCompanionAdapter::MergeFunction::extractValues(
        char **groups,
        int32_t numGroups,
        VectorPtr *result) {
        fn_->extractAccumulators(groups, numGroups, result);
    }

    void AggregateCompanionAdapter::MergeExtractFunction::extractValues(
        char **groups,
        int32_t numGroups,
        VectorPtr *result) {
        fn_->extractValues(groups, numGroups, result);
    }

    int32_t AggregateCompanionAdapter::ExtractFunction::setOffset() const {
        const int32_t rowSizeOffset = bits::nbytes(1);
        // Tracked row size takes an uint32_t slot. Accumulator starts after the
        // row-size slot.
        int32_t offset = rowSizeOffset + sizeof(uint32_t);
        offset = bits::roundUp(offset, fn_->accumulatorAlignmentSize());
        fn_->setOffsets(
            offset,
            RowContainer::nullByte(0),
            RowContainer::nullMask(0),
            RowContainer::initializedByte(0),
            RowContainer::initializedMask(0),
            rowSizeOffset);
        return offset;
    }

    char **AggregateCompanionAdapter::ExtractFunction::allocateGroups(
        memory::AllocationPool &allocationPool,
        const SelectivityVector &rows,
        uint64_t offsetInGroup) const {
        auto *groups = reinterpret_cast<char **>(
            allocationPool.allocateFixed(sizeof(char *) * rows.end()));

        auto size = fn_->accumulatorFixedWidthSize();
        auto alignment = fn_->accumulatorAlignmentSize();
        rows.applyToSelected([&](auto row) {
            groups[row] = allocationPool.allocateFixed(size + offsetInGroup, alignment);
        });
        return groups;
    }

    std::tuple<vector_size_t, BufferPtr>
    AggregateCompanionAdapter::ExtractFunction::compactGroups(
        memory::MemoryPool *pool,
        const SelectivityVector &rows,
        char **groups) const {
        BufferPtr indices = allocate_indices(rows.end(), pool);
        auto *rawIndices = indices->asMutable<vector_size_t>();
        vector_size_t count = 0;
        rows.applyToSelected([&](auto row) {
            if (count < row) {
                groups[count] = groups[row];
            }
            rawIndices[row] = count;
            ++count;
        });
        return std::make_tuple(count, indices);
    }

    void AggregateCompanionAdapter::ExtractFunction::apply(
        const SelectivityVector &rows,
        std::vector<VectorPtr> &args,
        const TypePtr &outputType,
        exec::EvalCtx &context,
        VectorPtr &result) const {
        // Set up data members of fn_.
        HashStringAllocator stringAllocator{context.pool()};
        memory::AllocationPool allocationPool{context.pool()};
        fn_->setAllocator(&stringAllocator);

        auto offset = setOffset();
        char **groups = allocateGroups(allocationPool, rows, offset);

        // Perform per-row aggregation.
        std::vector<vector_size_t> allSelectedRange;
        rows.applyToSelected([&](auto row) { allSelectedRange.push_back(row); });
        fn_->initializeNewGroups(groups, allSelectedRange);
        fn_->enableValidateIntermediateInputs();
        fn_->addIntermediateResults(groups, rows, args, false);

        auto localResult = BaseVector::create(outputType, rows.end(), context.pool());
        const auto &[groupCount, rowsToGroupsIndices] =
                compactGroups(context.pool(), rows, groups);
        fn_->extractValues(groups, groupCount, &localResult);
        localResult = BaseVector::wrap_in_dictionary(
            nullptr, rowsToGroupsIndices, rows.end(), localResult);
        context.moveOrCopyResult(localResult, rows, result);

        if (fn_->accumulatorUsesExternalMemory()) {
            fn_->destroy(melon::Range(groups, groupCount));
        }
    }

    bool CompanionFunctionsRegistrar::registerPartialFunction(
        const std::string &name,
        const std::vector<AggregateFunctionSignaturePtr> &signatures,
        const AggregateFunctionMetadata &metadata,
        bool overwrite) {
        auto partialSignatures =
                CompanionSignatures::partialFunctionSignatures(signatures);
        if (partialSignatures.empty()) {
            return false;
        }

        return exec::registerAggregateFunction(
                    CompanionSignatures::partialFunctionName(name),
                    std::move(partialSignatures),
                    [name](
                core::AggregationNode::Step step,
                const std::vector<TypePtr> &argTypes,
                const TypePtr &resultType,
                const core::QueryConfig &config)
                -> std::unique_ptr<Aggregate> {
                        if (auto func = getAggregateFunctionEntry(name)) {
                            if (!exec::isRawInput(step)) {
                                step = core::AggregationNode::Step::kIntermediate;
                            }
                            auto fn = func->factory(step, argTypes, resultType, config);
                            POLLUX_CHECK_NOT_NULL(fn);
                            return std::make_unique<
                                AggregateCompanionAdapter::PartialFunction>(
                                std::move(fn), resultType);
                        }
                        POLLUX_FAIL(
                            "Original aggregation function {} not found: {}",
                            name,
                            CompanionSignatures::partialFunctionName(name));
                    },
                    metadata,
                    /*registerCompanionFunctions*/ false,
                    overwrite)
                .mainFunction;
    }

    bool CompanionFunctionsRegistrar::registerMergeFunction(
        const std::string &name,
        const std::vector<AggregateFunctionSignaturePtr> &signatures,
        const AggregateFunctionMetadata &metadata,
        bool overwrite) {
        auto mergeSignatures =
                CompanionSignatures::mergeFunctionSignatures(signatures);
        if (mergeSignatures.empty()) {
            return false;
        }

        return exec::registerAggregateFunction(
                    CompanionSignatures::mergeFunctionName(name),
                    std::move(mergeSignatures),
                    [name](
                core::AggregationNode::Step /*step*/,
                const std::vector<TypePtr> &argTypes,
                const TypePtr &resultType,
                const core::QueryConfig &config)
                -> std::unique_ptr<Aggregate> {
                        if (auto func = getAggregateFunctionEntry(name)) {
                            auto fn = func->factory(
                                core::AggregationNode::Step::kIntermediate,
                                argTypes,
                                resultType,
                                config);
                            POLLUX_CHECK_NOT_NULL(fn);
                            return std::make_unique<
                                AggregateCompanionAdapter::MergeFunction>(
                                std::move(fn), resultType);
                        }
                        POLLUX_FAIL(
                            "Original aggregation function {} not found: {}",
                            name,
                            CompanionSignatures::mergeFunctionName(name));
                    },
                    metadata,
                    /*registerCompanionFunctions*/ false,
                    overwrite)
                .mainFunction;
    }

    bool registerMergeExtractFunctionInternal(
        const std::string &name,
        const std::string &mergeExtractFunctionName,
        const std::vector<std::shared_ptr<AggregateFunctionSignature> > &
        mergeExtractSignatures,
        const AggregateFunctionMetadata &metadata,
        bool overwrite) {
        return exec::registerAggregateFunction(
                    mergeExtractFunctionName,
                    std::move(mergeExtractSignatures),
                    [name, mergeExtractFunctionName](
                core::AggregationNode::Step /*step*/,
                const std::vector<TypePtr> &argTypes,
                const TypePtr &resultType,
                const core::QueryConfig &config)
                -> std::unique_ptr<Aggregate> {
                        const auto &[originalResultType, _] =
                                resolveAggregateFunction(mergeExtractFunctionName, argTypes);
                        if (!originalResultType) {
                            // TODO: limitation -- result type must be resolvable given
                            // intermediate type of the original UDAF.
                            POLLUX_UNREACHABLE(
                                "Signatures whose result types are not resolvable given intermediate types should have been excluded.")
                            ;
                        }

                        if (auto func = getAggregateFunctionEntry(name)) {
                            auto fn = func->factory(
                                core::AggregationNode::Step::kFinal,
                                argTypes,
                                originalResultType,
                                config);
                            POLLUX_CHECK_NOT_NULL(fn);
                            return std::make_unique<
                                AggregateCompanionAdapter::MergeExtractFunction>(
                                std::move(fn), resultType);
                        }
                        POLLUX_FAIL(
                            "Original aggregation function {} not found: {}",
                            name,
                            mergeExtractFunctionName);
                    },
                    metadata,
                    /*registerCompanionFunctions*/ false,
                    overwrite)
                .mainFunction;
    }

    bool CompanionFunctionsRegistrar::registerMergeExtractFunctionWithSuffix(
        const std::string &name,
        const std::vector<AggregateFunctionSignaturePtr> &signatures,
        const AggregateFunctionMetadata &metadata,
        bool overwrite) {
        auto groupedSignatures =
                CompanionSignatures::groupSignaturesByReturnType(signatures);
        bool registered = false;
        for (const auto &[type, signatureGroup]: groupedSignatures) {
            auto mergeExtractSignatures =
                    CompanionSignatures::mergeExtractFunctionSignatures(signatureGroup);
            if (mergeExtractSignatures.empty()) {
                continue;
            }

            auto mergeExtractFunctionName =
                    CompanionSignatures::mergeExtractFunctionNameWithSuffix(name, type);

            registered |= registerMergeExtractFunctionInternal(
                name,
                mergeExtractFunctionName,
                std::move(mergeExtractSignatures),
                metadata,
                overwrite);
        }
        return registered;
    }

    bool CompanionFunctionsRegistrar::registerMergeExtractFunction(
        const std::string &name,
        const std::vector<AggregateFunctionSignaturePtr> &signatures,
        const AggregateFunctionMetadata &metadata,
        bool overwrite) {
        if (CompanionSignatures::hasSameIntermediateTypesAcrossSignatures(
            signatures)) {
            return registerMergeExtractFunctionWithSuffix(
                name, signatures, metadata, overwrite);
        }

        auto mergeExtractSignatures =
                CompanionSignatures::mergeExtractFunctionSignatures(signatures);
        if (mergeExtractSignatures.empty()) {
            return false;
        }

        auto mergeExtractFunctionName =
                CompanionSignatures::mergeExtractFunctionName(name);
        return registerMergeExtractFunctionInternal(
            name,
            mergeExtractFunctionName,
            std::move(mergeExtractSignatures),
            metadata,
            overwrite);
    }

    VectorFunctionFactory getVectorFunctionFactory(
        const std::string &originalName) {
        return [originalName](
            const std::string &name,
            const std::vector<VectorFunctionArg> &inputArgs,
            const core::QueryConfig &config)
            -> std::shared_ptr<VectorFunction> {
            std::vector<TypePtr> argTypes{inputArgs.size()};
            std::transform(
                inputArgs.begin(),
                inputArgs.end(),
                argTypes.begin(),
                [](auto inputArg) { return inputArg.type; });

            auto resultType = resolveVectorFunction(name, argTypes);
            if (!resultType) {
                // TODO: limitation -- result type must be resolvable given
                // intermediate type of the original UDAF.
                POLLUX_UNREACHABLE(
                    "Signatures whose result types are not resolvable given intermediate types should have been excluded.")
                ;
            }

            if (auto func = getAggregateFunctionEntry(originalName)) {
                auto fn = func->factory(
                    core::AggregationNode::Step::kFinal, argTypes, resultType, config);
                POLLUX_CHECK_NOT_NULL(fn);
                return std::make_shared<AggregateCompanionAdapter::ExtractFunction>(
                    std::move(fn));
            }
            POLLUX_FAIL(
                "Original aggregation function {} not found: {}", originalName, name);
        };
    }

    bool CompanionFunctionsRegistrar::registerExtractFunctionWithSuffix(
        const std::string &originalName,
        const std::vector<AggregateFunctionSignaturePtr> &signatures,
        bool overwrite) {
        auto groupedSignatures =
                CompanionSignatures::groupSignaturesByReturnType(signatures);
        bool registered = false;
        for (const auto &[type, signatureGroup]: groupedSignatures) {
            auto extractSignatures =
                    CompanionSignatures::extractFunctionSignatures(signatureGroup);
            if (extractSignatures.empty()) {
                continue;
            }

            auto factory = getVectorFunctionFactory(originalName);
            registered |= exec::registerStatefulVectorFunction(
                CompanionSignatures::extractFunctionNameWithSuffix(originalName, type),
                std::move(extractSignatures),
                std::move(factory),
                exec::VectorFunctionMetadataBuilder()
                .defaultNullBehavior(false)
                .companionFunction(true)
                .build(),
                overwrite);
        }
        return registered;
    }

    bool CompanionFunctionsRegistrar::registerExtractFunction(
        const std::string &originalName,
        const std::vector<AggregateFunctionSignaturePtr> &signatures,
        bool overwrite) {
        if (CompanionSignatures::hasSameIntermediateTypesAcrossSignatures(
            signatures)) {
            return registerExtractFunctionWithSuffix(
                originalName, signatures, overwrite);
        }

        auto extractSignatures =
                CompanionSignatures::extractFunctionSignatures(signatures);
        if (extractSignatures.empty()) {
            return false;
        }

        auto factory = getVectorFunctionFactory(originalName);
        return exec::registerStatefulVectorFunction(
            CompanionSignatures::extractFunctionName(originalName),
            std::move(extractSignatures),
            std::move(factory),
            exec::VectorFunctionMetadataBuilder()
            .defaultNullBehavior(false)
            .companionFunction(true)
            .build(),
            overwrite);
    }
} // namespace kumo::pollux::exec
